//
// Created by ZhaoXiaoFei on 2022/7/21.
//
#include "initial_sfm.hpp"

void GlobalSfm::triangulation(Eigen::Matrix<double, 3, 4>& pose1, Eigen::Vector2d& point1,
                              Eigen::Matrix<double, 3, 4>& pose2, Eigen::Vector2d& point2, Eigen::Vector3d& point_3d){
    Eigen::Matrix4d A = Eigen::Matrix4d::Zero();
    A.row(0) = pose1.row(2) * point1(0) - pose1.row(0);
    A.row(1) = pose1.row(2) * point1(1) - pose1.row(1);
    A.row(2) = pose2.row(2) * point2(0) - pose2.row(0);
    A.row(3) = pose2.row(2) * point2(1) - pose2.row(1);

    Eigen::Vector4d res = A.jacobiSvd(Eigen::ComputeFullV).matrixV().block<4, 1>(0, 3);

    point_3d(0) = res(0 ) / res(3);
    point_3d(1) = res(1 ) / res(3);
    point_3d(2) = res(2 ) / res(3);
}

void GlobalSfm::triangulationTwoFrame(int frame_i, Eigen::Matrix<double, 3, 4> &pose1, int frame_j,
                                      Eigen::Matrix<double, 3, 4> &pose2, std::vector<SfmFeature>& sfmVec) {
    assert(frame_i != frame_j);
    for(int i = 0; i < sfmVec.size(); i++){
        if(sfmVec[i].state == true){
            continue;
        }
        bool has_i = false;
        bool has_j = false;
        Eigen::Vector2d point_i, point_j;
        for(int j = 0; j < sfmVec[i].observation.size(); j++){
            if(sfmVec[i].observation[j].first == frame_i){
                has_i = true;
                point_i = sfmVec[i].observation[j].second;
            }
            if(sfmVec[i].observation[j].first == frame_j){
                has_j = true;
                point_j = sfmVec[i].observation[j].second;
            }
        }
        if(has_i && has_j){
            Eigen::Vector3d point_3d;
            triangulation(pose1, point_i, pose2, point_j, point_3d);
            sfmVec[i].state = true;
            sfmVec[i].position[0] = point_3d[0];
            sfmVec[i].position[1] = point_3d[1];
            sfmVec[i].position[2] = point_3d[2];
            //LOG(INFO) << "trangulated : " << frame_j << "  3d point : "  << i << "  " << point_3d.transpose() << std::endl;
        }
    }

}

bool GlobalSfm::solveByPnP(Eigen::Matrix3d &R, Eigen::Vector3d &T, int &k, std::vector<SfmFeature> &sfmVec) {
    std::vector<cv::Point2f> points_2d;
    std::vector<cv::Point3f> points_3d;
    for(int i = 0; i < sfmVec.size(); i++){
        if(sfmVec[i].state != true){
            continue;
        }
        for(int j = 0; j < sfmVec[i].observation.size(); j++){
            if(sfmVec[i].observation[j].first == k){
                cv::Point2f point_2d(sfmVec[i].observation[j].second[0], sfmVec[i].observation[j].second[1]);
                points_2d.push_back(point_2d);
                cv::Point3f point_3d(sfmVec[i].position[0], sfmVec[i].position[1], sfmVec[i].position[2]);
                points_3d.push_back(point_3d);
                break;
            }
        }
    }
    if (int(points_2d.size()) < 15)
    {
        LOG(INFO) << "points_2d.size(): " << points_2d.size() << std::endl;
        LOG(INFO) << "unstable features tracking, please slowly move you device!" << std::endl;
        if (int(points_2d.size()) < 10)
            return false;
    }
//    LOG(INFO) << "R1111: " << R << std::endl;
    //cv::Mat K = cv::Mat::eye(3, 3, CV_64F);
    cv::Mat D, rvec, t, tempR;
    cv::eigen2cv(R, tempR);
    cv::Rodrigues(tempR, rvec);
    cv::eigen2cv(T, t);
    cv::Mat K = (cv::Mat_<double>(3, 3) << 1, 0, 0, 0, 1, 0, 0, 0, 1);
    bool pnp_res;
    pnp_res = cv::solvePnP(points_3d, points_2d, K, D, rvec, t, 1);
    if(!pnp_res){
        return false;
    }
    //LOG(INFO) << "r: " << r << std::endl;
    cv::Rodrigues(rvec, tempR);
    //LOG(INFO) << "tempR: " << tempR << std::endl;
    cv::cv2eigen(tempR, R);
//    LOG(INFO) << "R2222: " << R << std::endl;
    cv::cv2eigen(t, T);
    return true;
}

bool GlobalSfm::construct(int frame_num, Eigen::Quaterniond *Q, Eigen::Vector3d *T, int l,
                          const Eigen::Matrix3d &relativeR, const Eigen::Vector3d &relativeT,
                          std::vector<SfmFeature> &sfmVec,
                          std::map<int, Eigen::Vector3d>& sfm_tracked_points) {
    //R_w_ck
    Q[l].w() = 1;
    Q[l].x() = 0;
    Q[l].y() = 0;
    Q[l].z() = 0;
    T[l].setZero();
    Q[frame_num - 1] = Q[l] * Eigen::Quaterniond(relativeR);
    T[frame_num - 1] = relativeT;

    // Q_c_w;
    double Q_c_w[frame_num][4];
    double T_c_w[frame_num][3];

    Eigen::Matrix<double, 3, 4> Pose[frame_num];
    Pose[l].block<3, 3>(0 ,0) = Q[l].toRotationMatrix().transpose();
    Pose[l].block<3, 1>(0 ,3) = -Q[l].toRotationMatrix().transpose() * T[l];

    Pose[frame_num - 1].block<3, 3>(0, 0) = Q[frame_num - 1].toRotationMatrix().transpose();
    Pose[frame_num - 1].block<3, 1>(0, 3) = -Q[frame_num - 1].toRotationMatrix().transpose() * T[frame_num - 1];

    for(int i = l; i < frame_num - 1; i++){
        if(i > l){
            Eigen::Matrix3d R_initial = Pose[i - 1].block<3, 3>(0, 0);
            Eigen::Vector3d T_initial = Pose[i - 1].block<3, 1>(0, 3);
            //LOG(INFO) << "R_initial1: " << R_initial << std::endl;
            if(!solveByPnP(R_initial, T_initial, i, sfmVec)){
                return false;
            }
            //LOG(INFO) << "R_initial2: " << R_initial << std::endl;
            Pose[i].block<3, 3>(0 ,0) = R_initial;
            Pose[i].block<3, 1>(0 ,3) = T_initial;
        }
        triangulationTwoFrame(i, Pose[i], frame_num - 1, Pose[frame_num - 1], sfmVec);
    }
    for(int i = l + 1; i < frame_num - 1; i++){
        triangulationTwoFrame(l, Pose[l], i, Pose[i], sfmVec);
    }

    for(int i = l - 1; i >= 0; i--){
        Eigen::Matrix3d R_initial = Pose[i + 1].block<3, 3>(0, 0);
        Eigen::Vector3d T_initial = Pose[i + 1].block<3, 1>(0, 3);
        if(!solveByPnP(R_initial, T_initial, i, sfmVec)){
            return false;
        }
        Pose[i].block<3, 3>(0 ,0) = R_initial;
        Pose[i].block<3, 1>(0 ,3) = T_initial;

        triangulationTwoFrame(i, Pose[i], l, Pose[l], sfmVec);
    }

    for(int i = 0; i < sfmVec.size(); i++){
        if(sfmVec[i].state){
            continue;
        }
        if(sfmVec[i].observation.size() >= 2){
            Eigen::Vector2d point1, point2;
            int frame_i = sfmVec[i].observation[0].first;
            point1 = sfmVec[i].observation[0].second;
            int frame_j = sfmVec[i].observation.back().first;
            point2 = sfmVec[i].observation.back().second;

            Eigen::Vector3d point_3d;
            triangulation(Pose[frame_i], point1, Pose[frame_j], point2, point_3d);
            sfmVec[i].state = true;
            sfmVec[i].position[0] = point_3d(0);
            sfmVec[i].position[1] = point_3d(1);
            sfmVec[i].position[2] = point_3d(2);
        }
    }

//    for(int i = 0; i < frame_num; i++){
//        LOG(INFO) << "Pose[" << i << "]: " << Pose[i] << std::endl;
//    }

    ceres::Problem problem;
    ceres::LocalParameterization* localParameterization = new ceres::QuaternionParameterization();
    for(int i = 0; i < frame_num; i++){
        Eigen::Quaterniond temp(Pose[i].block<3, 3>(0, 0));
        Q_c_w[i][0] = temp.w();
        Q_c_w[i][1] = temp.x();
        Q_c_w[i][2] = temp.y();
        Q_c_w[i][3] = temp.z();
        T_c_w[i][0] = Pose[i](0, 3);
        T_c_w[i][1] = Pose[i](1, 3);
        T_c_w[i][2] = Pose[i](2, 3);
        problem.AddParameterBlock(Q_c_w[i], 4, localParameterization);
        problem.AddParameterBlock(T_c_w[i], 3);
        if(i == l){
            problem.SetParameterBlockConstant(Q_c_w[i]);
        }
        if(i == l || i == frame_num - 1){
            problem.SetParameterBlockConstant(T_c_w[i]);
        }
    }
    ceres::LossFunction* lossFunction = new ceres::HuberLoss(0.1);
    for(int i = 0; i < sfmVec.size(); i++){
        if(sfmVec[i].state != true){
            continue;
        }
        for(int j = 0; j < sfmVec[i].observation.size(); j++){
            int k = sfmVec[i].observation[j].first;
            ceres::CostFunction* costFunction = ReprojectionError3D::Create(sfmVec[i].observation[j].second.x(),
                                                                          sfmVec[i].observation[j].second.y());
            problem.AddResidualBlock(costFunction, nullptr, Q_c_w[k], T_c_w[k],
                                     sfmVec[i].position);
        }
    }
    ceres::Solver::Options options;
    options.linear_solver_type = ceres::DENSE_SCHUR;
    options.minimizer_progress_to_stdout = true;
    options.max_solver_time_in_seconds = 0.2;
    ceres::Solver::Summary summary;
    ceres::Solve(options, &problem, &summary);
    std::cout << summary.BriefReport() << "\n";
    if (summary.termination_type == ceres::CONVERGENCE || summary.final_cost < 5e-03)
    {
        LOG(INFO)  << "vision only BA converge" << std::endl;
    }
    else
    {
        LOG(INFO) << "vision only BA not converge " << std::endl;
        return false;
    }
    for(int i = 0; i < frame_num; i++){
        Eigen::Quaterniond temp;
        temp.w() = Q_c_w[i][0];
        temp.x() = Q_c_w[i][1];
        temp.y() = Q_c_w[i][2];
        temp.z() = Q_c_w[i][3];
        Q[i] = temp.inverse();

        Eigen::Vector3d temp_T;
        temp_T(0) = T_c_w[i][0];
        temp_T(1) = T_c_w[i][1];
        temp_T(2) = T_c_w[i][2];
        T[i] = -Q[i].toRotationMatrix() * temp_T;
    }
    for(int i = 0; i < sfmVec.size(); i++){
        if(sfmVec[i].state){
            sfm_tracked_points.insert(std::make_pair(sfmVec[i].id, Eigen::Vector3d(sfmVec[i].position[0],
                                                                                   sfmVec[i].position[1],
                                                                                   sfmVec[i].position[2])));
//            LOG(INFO) << Eigen::Vector3d(sfmVec[i].position[0],
//                                         sfmVec[i].position[1],
//                                         sfmVec[i].position[2]).transpose() << std::endl;
        }

    }

//    LOG(INFO) << "SFM detile!!!" << std::endl;
//    for(int i = 0; i < frame_num; i++){
//        LOG(INFO) << "Q: " << i << " " << Q[i] << std::endl;
//        LOG(INFO) << "T: " << i << " " << T[i].transpose() << std::endl;
//    }
//    LOG(INFO) << "~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~`" << std::endl;


    return true;
}